#include "ns3/log.h"
#include "ns3/enum.h"
#include "ns3/uinteger.h"
#include "ns3/double.h"
#include "ns3/simulator.h"
#include "ns3/abort.h"
#include "ns3/drop-tail-queue.h"
#include "C-pie-queue.h"
#include "CuckooHash.h"


// Edit by PassLink.cn
namespace ns3
{

  // ****************************添加代码************************************//
  // 创建一个 CuckooHash 对象，假设初始大小为1000
  CuckooHash cuckooHash(1000);
  // ****************************添加代码************************************//
  
  NS_LOG_COMPONENT_DEFINE("cPieQueueDisc");

  NS_OBJECT_ENSURE_REGISTERED(cPieQueueDisc);

  TypeId cPieQueueDisc::GetTypeId(void)
  {
    static TypeId tid = TypeId("ns3::cPieQueueDisc")
                            .SetParent<QueueDisc>()
                            .SetGroupName("TrafficControl")
                            .AddConstructor<cPieQueueDisc>()
                            // .AddAttribute("CuckooHashSize",
                            //               "Size of cuckooHash ",
                            //               UintegerValue(1000),
                            //               MakeUintegerAccessor(&cPieQueueDisc::m_cuckooHashSize),
                            //               MakeUintegerChecker<uint32_t>())
                            .AddAttribute("MeanPktSize",
                                          "Average of packet size",
                                          UintegerValue(1000),
                                          MakeUintegerAccessor(&cPieQueueDisc::m_meanPktSize),
                                          MakeUintegerChecker<uint32_t>())
                            .AddAttribute("A",
                                          "Value of alpha",
                                          DoubleValue(0.125),
                                          MakeDoubleAccessor(&cPieQueueDisc::m_a),
                                          MakeDoubleChecker<double>())
                            .AddAttribute("B",
                                          "Value of beta",
                                          DoubleValue(1.25),
                                          MakeDoubleAccessor(&cPieQueueDisc::m_b),
                                          MakeDoubleChecker<double>())
                            .AddAttribute("Tupdate",
                                          "Time period to calculate drop probability",
                                          TimeValue(MilliSeconds(15)),
                                          MakeTimeAccessor(&cPieQueueDisc::m_tUpdate),
                                          MakeTimeChecker())
                            .AddAttribute("Supdate",
                                          "Start time of the update timer",
                                          TimeValue(Seconds(0)),
                                          MakeTimeAccessor(&cPieQueueDisc::m_sUpdate),
                                          MakeTimeChecker())
                            .AddAttribute("MaxSize",
                                          "The maximum number of packets accepted by this queue disc",
                                          QueueSizeValue(QueueSize("25p")),
                                          MakeQueueSizeAccessor(&QueueDisc::SetMaxSize,
                                                                &QueueDisc::GetMaxSize),
                                          MakeQueueSizeChecker())
                            .AddAttribute("DequeueThreshold",
                                          "Minimum queue size in bytes before dequeue rate is measured",
                                          UintegerValue(16384),
                                          MakeUintegerAccessor(&cPieQueueDisc::m_dqThreshold),
                                          MakeUintegerChecker<uint32_t>())
                            .AddAttribute("QueueDelayReference",
                                          "Desired queue delay",
                                          TimeValue(MilliSeconds(15)),
                                          MakeTimeAccessor(&cPieQueueDisc::m_qDelayRef),
                                          MakeTimeChecker())
                            .AddAttribute("MaxBurstAllowance",
                                          "Current max burst allowance before random drop",
                                          TimeValue(MilliSeconds(15)),
                                          MakeTimeAccessor(&cPieQueueDisc::m_maxBurst),
                                          MakeTimeChecker())
                            .AddAttribute("UseDequeueRateEstimator",
                                          "Enable/Disable usage of Dequeue Rate Estimator",
                                          BooleanValue(false),
                                          MakeBooleanAccessor(&cPieQueueDisc::m_useDqRateEstimator),
                                          MakeBooleanChecker())
                            .AddAttribute("UseCapDropAdjustment",
                                          "Enable/Disable Cap Drop Adjustment feature mentioned in RFC 8033",
                                          BooleanValue(true),
                                          MakeBooleanAccessor(&cPieQueueDisc::m_isCapDropAdjustment),
                                          MakeBooleanChecker())
                            .AddAttribute("UseEcn",
                                          "True to use ECN (packets are marked instead of being dropped)",
                                          BooleanValue(false),
                                          MakeBooleanAccessor(&cPieQueueDisc::m_useEcn),
                                          MakeBooleanChecker())
                            .AddAttribute("MarkEcnThreshold",
                                          "ECN marking threshold (RFC 8033 suggests 0.1 (i.e., 10%) default)",
                                          DoubleValue(0.1),
                                          MakeDoubleAccessor(&cPieQueueDisc::m_markEcnTh),
                                          MakeDoubleChecker<double>(0, 1))
                            .AddAttribute("UseDerandomization",
                                          "Enable/Disable Derandomization feature mentioned in RFC 8033",
                                          BooleanValue(false),
                                          MakeBooleanAccessor(&cPieQueueDisc::m_useDerandomization),
                                          MakeBooleanChecker())
                            .AddAttribute("ActiveThreshold",
                                          "Threshold for activating PIE (disabled by default)",
                                          TimeValue(Time::Max()),
                                          MakeTimeAccessor(&cPieQueueDisc::m_activeThreshold),
                                          MakeTimeChecker())
                            .AddAttribute("CeThreshold",
                                          "The FqPie CE threshold for marking packets",
                                          TimeValue(Time::Max()),
                                          MakeTimeAccessor(&cPieQueueDisc::m_ceThreshold),
                                          MakeTimeChecker())
                            .AddAttribute("UseL4s",
                                          "True to use L4S (only ECT1 packets are marked at CE threshold)",
                                          BooleanValue(false),
                                          MakeBooleanAccessor(&cPieQueueDisc::m_useL4s),
                                          MakeBooleanChecker());

    return tid;
  }




  // 初始化
  cPieQueueDisc::cPieQueueDisc()
      : QueueDisc(QueueDiscSizePolicy::SINGLE_INTERNAL_QUEUE)
  { 
    NS_LOG_FUNCTION(this);
    m_uv = CreateObject<UniformRandomVariable>();
    m_rtrsEvent = Simulator::Schedule(m_sUpdate, &cPieQueueDisc::CalculateP, this);
  }

  // 销毁工作
  cPieQueueDisc::~cPieQueueDisc()
  {
    NS_LOG_FUNCTION(this);
  }

  // 销毁工作
  void
  cPieQueueDisc::DoDispose(void)
  {
    NS_LOG_FUNCTION(this);
    m_uv = 0;
    m_rtrsEvent.Cancel();
    QueueDisc::DoDispose();
  }

  // 获取队列延迟
  Time
  cPieQueueDisc::GetQueueDelay(void)
  {
    return m_qDelay;
  }

  // 用于为PieQueueDisc对象分配流（Streams）
  int64_t
  cPieQueueDisc::AssignStreams(int64_t stream)
  {
    NS_LOG_FUNCTION(this << stream);
    m_uv->SetStream(stream);
    return 1;
  }

  // 入队操作
  bool
  cPieQueueDisc::DoEnqueue(Ptr<QueueDiscItem> item)
  {
    Ptr<Ipv4QueueDiscItem> ipv4Item = DynamicCast<Ipv4QueueDiscItem>(item);
    string ipv4ItemHash = to_string(ipv4Item->Hash(0));

    NS_LOG_FUNCTION(this << item);
    // 调用GetCurrentSize()函数获取当前队列的大小，并将其存储在nQueued变量中。
    QueueSize nQueued = GetCurrentSize();
    // If L4S is enabled, then check if the packet is ECT1, and if it is then set isEct true
    bool isEct1 = false;
    // 根据队列的状态和数据包的特性进行处理。
    // 这段代码检查是否启用了L4S（低延迟、低丢包）功能，并且检查数据包是否设置了ECT1标志。
    // 如果是，则设置isEct1为true，并记录日志。
    if (item && m_useL4s)
    {
      uint8_t tosByte = 0;
      if (item->GetUint8Value(QueueItem::IP_DSFIELD, tosByte) && (((tosByte & 0x3) == 1) || (tosByte & 0x3) == 3))
      {
        if ((tosByte & 0x3) == 1)
        {
          NS_LOG_DEBUG("Enqueueing ECT1 packet " << static_cast<uint16_t>(tosByte & 0x3));
        }
        else
        {
          NS_LOG_DEBUG("Enqueueing CE packet " << static_cast<uint16_t>(tosByte & 0x3));
        }
        isEct1 = true;
      }
    }

    // 检查加入数据包后队列是否超出了最大大小。如果超出了，
    // 根据队列的状态进行处理，这里是强制性地丢弃数据包，并返回false。
    if (nQueued + item > GetMaxSize())
    {
      // Drops due to queue limit: reactive
      DropBeforeEnqueue(item, FORCED_DROP);
      m_accuProb = 0;
      return false;
    }

    // isEct1 will be true only if L4S enabled as well as the packet is ECT1.
    // If L4S is enabled and packet is ECT1 then directly enqueue the packet.
    // 这段代码处理早期概率丢包，根据队列的状态和数据包的特性进行处理。
    else if ((m_activeThreshold == Time::Max() || m_active) && !isEct1 && DropEarly(item, nQueued.GetValue()))
    {
      if (!m_useEcn || m_dropProb >= m_markEcnTh || !Mark(item, UNFORCED_MARK))
      {
        // Early probability drop: proactive
        DropBeforeEnqueue(item, UNFORCED_DROP);
        m_accuProb = 0;
        return false;
      }
    }

    

    // 这段代码检查队列延迟是否超过了特定阈值，如果是，则启用PIE算法。
    // If the queue is over a certain threshold, Turn ON PIE
    // ****************************添加代码************************************//



    if (m_activeThreshold != Time::Max() && !m_active && m_qDelay >= (m_activeThreshold / 9) && cuckooHash.isFloodStreamflow(ipv4ItemHash))
      { 
        // std::cout << "m_qDelay= "<< m_qDelay << std::endl;
        // std::cout << "m_activeThreshold / 9 = "<< m_activeThreshold / 9 << std::endl;
        // std::cout << "ipv4ItemHash:"<<ipv4ItemHash << std::endl;
 
        // Early probability drop: proactive
        DropBeforeEnqueue(item, UNFORCED_DROP);
        m_accuProb = 0;
        return false;
      }

    // ****************************添加代码************************************//
    if (m_activeThreshold != Time::Max() && !m_active && m_qDelay < m_activeThreshold / 9)
    {
      m_active = true;
      m_qDelayOld = Seconds(0);
      m_dropProb = 0;
      m_inMeasurement = true;
      m_dqCount = 0;
      m_avgDqRate = 0;
      m_burstAllowance = m_maxBurst;
      m_accuProb = 0;
      m_dqStart = Now();
    }



    // No drop 这行代码将数据包加入内部队列，并将结果存储在retval变量中。
    bool retval = GetInternalQueue(0)->Enqueue(item);
    
    // ****************************添加代码************************************//
    // If queue has been Idle for a while, Turn OFF PIE
    // Reset Counters when accessing the queue after some idle period if PIE was active before
    // 这段代码检查队列是否长时间处于空闲状态，如果是，则禁用PIE算法。
    // if (m_activeThreshold != Time::Max() && m_dropProb == 0 && m_qDelayOld.GetMilliSeconds() == 0 && m_qDelay.GetMilliSeconds() == 0)
    // {
    //   m_active = false;
    //   m_inMeasurement = false;
    // }
      m_active = false;
      m_inMeasurement = false;
    // If Queue::Enqueue fails, QueueDisc::DropBeforeEnqueue is called by the
    // internal queue because QueueDisc::AddInternalQueue sets the trace callback

    NS_LOG_LOGIC("\t bytesInQueue  " << GetInternalQueue(0)->GetNBytes());
    NS_LOG_LOGIC("\t packetsInQueue  " << GetInternalQueue(0)->GetNPackets());

    // ****************************添加代码************************************//
    // 获取源IP、目的IP、源端口、目的端口和传输层协议等信息
    // std::cout << "ipv4ItemHash: "<< ipv4ItemHash << std::endl;
    cuckooHash.insert_into(ipv4ItemHash, 1);
    // ****************************添加代码************************************//
    return retval;
  }

  // 初始化变量
  void
  cPieQueueDisc::InitializeParams(void)
  {
    // Initially queue is empty so variables are initialize to zero except m_dqCount
    m_inMeasurement = false;
    m_dqCount = DQCOUNT_INVALID;
    m_dropProb = 0;
    m_avgDqRate = 0.0;
    m_dqStart = Seconds(0);
    m_burstState = NO_BURST;
    m_qDelayOld = Seconds(0);
    m_accuProb = 0.0;
    m_active = false;

    // 初始化定时清理阈值为2秒
    m_clearThreshold = Seconds(2);
    // 上次清理时间
    m_lastClearTime = Now();
  }

  // 根据早期丢包算法的各种情况进行处理。
  bool cPieQueueDisc::DropEarly(Ptr<QueueDiscItem> item, uint32_t qSize)
  {
    NS_LOG_FUNCTION(this << item << qSize);
    // 如果仍然有爆发允许值（burst_allowance），则跳过早期丢包。
    if (m_burstAllowance.GetSeconds() > 0)
    {
      // If there is still burst_allowance left, skip random early drop.
      return false;
    }

    // 如果当前没有爆发，则将爆发状态设置为“IN_BURST_PROTECTING”，并将爆发允许值设置为最大爆发值。
    if (m_burstState == NO_BURST)
    {
      m_burstState = IN_BURST_PROTECTING;
      m_burstAllowance = m_maxBurst;
    }
    // 计算丢包概率p。如果队列大小的单位是字节，则根据数据包大小和平均包大小来计算丢包概率。
    double p = m_dropProb;

    uint32_t packetSize = item->GetSize();

    if (GetMaxSize().GetUnit() == QueueSizeUnit::BYTES)
    {
      p = p * packetSize / m_meanPktSize;
    }
    // // 根据队列延迟、丢包概率和队列大小来决定是否丢弃数据包。
    // Safeguard PIE to be work conserving (Section 4.1 of RFC 8033)
    if ((m_qDelayOld.GetSeconds() < (0.5 * m_qDelayRef.GetSeconds())) && (m_dropProb < 0.2))
    {
      return false;
    }
    else if (GetMaxSize().GetUnit() == QueueSizeUnit::BYTES && qSize <= 2 * m_meanPktSize)
    {
      return false;
    }
    else if (GetMaxSize().GetUnit() == QueueSizeUnit::PACKETS && qSize <= 2)
    {
      return false;
    }

    // 如果启用了随机性降低功能，根据累积丢包概率来决定是否丢包。
    if (m_useDerandomization)
    {
      if (m_dropProb == 0)
      {
        m_accuProb = 0;
      }
      m_accuProb += m_dropProb;
      if (m_accuProb < 0.85)
      {
        return false;
      }
      else if (m_accuProb >= 8.5)
      {
        return true;
      }
    }

    double u = m_uv->GetValue();
    if (u > p)
    {
      return false;
    }

    return true;
  }

  // 计算丢包概率
  void cPieQueueDisc::CalculateP()
  {
    NS_LOG_FUNCTION(this);
    Time qDelay;
    double p = 0.0;
    bool missingInitFlag = false;

    if (m_useDqRateEstimator)
    {
      if (m_avgDqRate > 0)
      {
        qDelay = Seconds(GetInternalQueue(0)->GetNBytes() / m_avgDqRate);
      }
      else
      {
        qDelay = Seconds(0);
        missingInitFlag = true;
      }
      m_qDelay = qDelay;
    }
    else
    {
      qDelay = m_qDelay;
    }
    NS_LOG_DEBUG("Queue delay while calculating probability: " << qDelay.GetMilliSeconds() << "ms");

    if (m_burstAllowance.GetSeconds() > 0)
    {
      m_dropProb = 0;
    }
    else
    {
      p = m_a * (qDelay.GetSeconds() - m_qDelayRef.GetSeconds()) + m_b * (qDelay.GetSeconds() - m_qDelayOld.GetSeconds());
      if (m_dropProb < 0.000001)
      {
        p /= 2048;
      }
      else if (m_dropProb < 0.00001)
      {
        p /= 512;
      }
      else if (m_dropProb < 0.0001)
      {
        p /= 128;
      }
      else if (m_dropProb < 0.001)
      {
        p /= 32;
      }
      else if (m_dropProb < 0.01)
      {
        p /= 8;
      }
      else if (m_dropProb < 0.1)
      {
        p /= 2;
      }
      else
      {
        // The pseudocode in Section 4.2 of RFC 8033 suggests to use this for
        // assignment of p, but this assignment causes build failure on Mac OS
        // p = p;
      }

      // Cap Drop Adjustment (Section 5.5 of RFC 8033)
      if (m_isCapDropAdjustment && (m_dropProb >= 0.1) && (p > 0.02))
      {
        p = 0.02;
      }
    }

    p += m_dropProb;

    // For non-linear drop in prob
    // Decay the drop probability exponentially (Section 4.2 of RFC 8033)
    if (qDelay.GetSeconds() == 0 && m_qDelayOld.GetSeconds() == 0)
    {
      p *= 0.98;
    }

    // bound the drop probability (Section 4.2 of RFC 8033)
    if (p < 0)
    {
      m_dropProb = 0;
    }
    else if (p > 1)
    {
      m_dropProb = 1;
    }
    else
    {
      m_dropProb = p;
    }

    // Section 4.4 #2
    if (m_burstAllowance < m_tUpdate)
    {
      m_burstAllowance = Seconds(0);
    }
    else
    {
      m_burstAllowance -= m_tUpdate;
    }

    uint32_t burstResetLimit = static_cast<uint32_t>(BURST_RESET_TIMEOUT / m_tUpdate.GetSeconds());
    if ((qDelay.GetSeconds() < 0.5 * m_qDelayRef.GetSeconds()) && (m_qDelayOld.GetSeconds() < (0.5 * m_qDelayRef.GetSeconds())) && (m_dropProb == 0) && !missingInitFlag)
    {
      m_dqCount = DQCOUNT_INVALID;
      m_avgDqRate = 0.0;
    }
    if ((qDelay.GetSeconds() < 0.5 * m_qDelayRef.GetSeconds()) && (m_qDelayOld.GetSeconds() < (0.5 * m_qDelayRef.GetSeconds())) && (m_dropProb == 0) && (m_burstAllowance.GetSeconds() == 0))
    {
      if (m_burstState == IN_BURST_PROTECTING)
      {
        m_burstState = IN_BURST;
        m_burstReset = 0;
      }
      else if (m_burstState == IN_BURST)
      {
        m_burstReset++;
        if (m_burstReset > burstResetLimit)
        {
          m_burstReset = 0;
          m_burstState = NO_BURST;
        }
      }
    }
    else if (m_burstState == IN_BURST)
    {
      m_burstReset = 0;
    }

    m_qDelayOld = qDelay;
    m_rtrsEvent = Simulator::Schedule(m_tUpdate, &cPieQueueDisc::CalculateP, this);
  }

  // 出队操作
  Ptr<QueueDiscItem>
  cPieQueueDisc::DoDequeue()
  {
    NS_LOG_FUNCTION(this);
    // 检查队列是否为空
    if (GetInternalQueue(0)->IsEmpty())
    {
      NS_LOG_LOGIC("Queue empty");
      return 0;
    }
    // 从队列中出队一个数据包
    Ptr<QueueDiscItem> item = GetInternalQueue(0)->Dequeue();

    // ****************************添加代码************************************//
    // 周期性清理队列内存
    Time difTime = Simulator::Now() - m_lastClearTime;
    if (difTime.GetSeconds() >= m_clearThreshold.GetSeconds())
    {
      cuckooHash.clearMemory();
    }
    // ****************************添加代码************************************//

    NS_ASSERT_MSG(item != nullptr, "Dequeue null, but internal queue not empty");

    // If L4S is enabled and packet is ECT1, then check if delay is greater
    // than CE threshold and if it is then mark the packet,
    // skip PIE steps, and return the item.
    // 如果启用了L4S，并且数据包是ECT1，则检查延迟是否超过CE阈值，如果是，则标记数据包，
    // 跳过PIE步骤，并返回数据包
    if (m_useL4s)
    {
      uint8_t tosByte = 0;
      if (item->GetUint8Value(QueueItem::IP_DSFIELD, tosByte) && (((tosByte & 0x3) == 1) || (tosByte & 0x3) == 3))
      {
        if ((tosByte & 0x3) == 1)
        {
          NS_LOG_DEBUG("ECT1 packet " << static_cast<uint16_t>(tosByte & 0x3));
        }
        else
        {
          NS_LOG_DEBUG("CE packet " << static_cast<uint16_t>(tosByte & 0x3));
        }
        if ((Now() - item->GetTimeStamp() > m_ceThreshold) && Mark(item, CE_THRESHOLD_EXCEEDED_MARK))
        {
          NS_LOG_LOGIC("Marking due to CeThreshold " << m_ceThreshold.GetSeconds());
        }
        return item;
      }
    }

    // if not in a measurement cycle and the queue has built up to dq_threshold,
    // start the measurement cycle
    // 如果启用了队列速率估算器，并且队列已经积累到dq_threshold，则开始测量周期
    if (m_useDqRateEstimator)
    {
      if ((GetInternalQueue(0)->GetNBytes() >= m_dqThreshold) && (!m_inMeasurement))
      {
        m_dqStart = Now();
        m_dqCount = 0;
        m_inMeasurement = true;
      }

      if (m_inMeasurement)
      {
        m_dqCount += item->GetSize();
        // 完成一个测量周期
        // done with a measurement cycle
        if (m_dqCount >= m_dqThreshold)
        {
          Time dqTime = Now() - m_dqStart;
          if (dqTime > Seconds(0))
          {
            if (m_avgDqRate == 0)
            {
              m_avgDqRate = m_dqCount / dqTime.GetSeconds();
            }
            else
            {
              m_avgDqRate = (0.5 * m_avgDqRate) + (0.5 * (m_dqCount / dqTime.GetSeconds()));
            }
          }
          NS_LOG_DEBUG("Average Dequeue Rate after Dequeue: " << m_avgDqRate);
          // 如果有足够的数据，则重新开始一个测量周期
          // restart a measurement cycle if there is enough data
          if (GetInternalQueue(0)->GetNBytes() > m_dqThreshold)
          {
            m_dqStart = Now();
            m_dqCount = 0;
            m_inMeasurement = true;
          }
          else
          {
            m_dqCount = 0;
            m_inMeasurement = false;
          }
        }
      }
    }
    else
    {
      // 计算队列延迟
      m_qDelay = Now() - item->GetTimeStamp();
      // 如果队列为空，则将队列延迟设置为0
      if (GetInternalQueue(0)->GetNBytes() == 0)
      {
        m_qDelay = Seconds(0);
      }
    }
    return item;
  }

  bool
  cPieQueueDisc::CheckConfig(void)
  {
    NS_LOG_FUNCTION(this);
    if (GetNQueueDiscClasses() > 0)
    {
      NS_LOG_ERROR("cPieQueueDisc cannot have classes");
      return false;
    }

    if (GetNPacketFilters() > 0)
    {
      NS_LOG_ERROR("cPieQueueDisc cannot have packet filters");
      return false;
    }

    if (GetNInternalQueues() == 0)
    {
      // add  a DropTail queue
      AddInternalQueue(CreateObjectWithAttributes<DropTailQueue<QueueDiscItem>>("MaxSize", QueueSizeValue(GetMaxSize())));
    }

    if (GetNInternalQueues() != 1)
    {
      NS_LOG_ERROR("cPieQueueDisc needs 1 internal queue");
      return false;
    }

    return true;
  }

} // namespace ns3
